module vert_coord_mod

	use const_mod
	use namelist_mod
	use mesh_mod

	implicit none

	private
	public vert_coord_init
	public vert_coord_final
	public hybrid_coord_calc_ph
	public hybrid_coord_calc_ph_lev
	public hyai
	public hybi


	real(r8), allocatable, dimension(:) :: hyai
  real(r8), allocatable, dimension(:) :: hybi
  real(r8), allocatable, dimension(:) :: hyam
  real(r8), allocatable, dimension(:) :: hybm

  real(r8) :: p0 = 1d5 ! Reference pressure (Pa)
contains

	subroutine vert_coord_init(num_lev)

		integer, intent(in) :: num_lev
		integer k 

		if (allocated(hyai)) deallocate(hyai)
		if (allocated(hybi)) deallocate(hyai)
		if (allocated(hyam)) deallocate(hyai)
		if (allocated(hybm)) deallocate(hyai)

		allocate(hyai(num_lev+1)); hyai = 0
		allocate(hybi(num_lev+1)); hybi = 0
		allocate(hyam(num_lev  )); hyam = 0
		allocate(hybm(num_lev  )); hybm = 0

		! select case(trim(test_case))
		! case ('dcmip1')
			call hybrid_coord_dcmip1(p0, hyai, hybi)
		! end select

		do k = 1, num_lev
			hyam(k) = 0.5 * (hyai(k) + hyai(k+1))
			hybm(k) = 0.5 * (hybi(k) + hybi(k+1))
		end do
		do k = 1, num_lev + 1
			mesh%half_lev(k) = hyai(k) + hybi(k)
		end do
		do k = 1, num_lev
			mesh%full_lev(k) = hyam(k) + hybm(k)
			mesh%full_dlev(k) = mesh%half_lev(k+1) - mesh%half_lev(k)
		end do
		
		do k = 2, num_lev
			mesh%half_dlev(k) = mesh%full_lev(k) - mesh%full_lev(k-1)
		end do
		mesh%half_dlev(1) = mesh%full_lev(1) - mesh%half_lev(1)
		mesh%half_dlev(num_lev+1) = mesh%half_lev(num_lev+1) - mesh%full_lev(num_lev)
		
	end subroutine vert_coord_init

	subroutine vert_coord_final()

		if (allocated(hyai)) deallocate(hyai)
		if (allocated(hybi)) deallocate(hybi)
		if (allocated(hyam)) deallocate(hyam)
		if (allocated(hybm)) deallocate(hybm)
		
	end subroutine vert_coord_final

	pure real(r8) function hybrid_coord_calc_ph(k, phs) result(res)

		integer, intent(in) :: k
		real(r8), intent(in) :: phs

		res = hyam(k) * p0 + hybm(k) * phs

	end function hybrid_coord_calc_ph

	pure real(r8) function hybrid_coord_calc_ph_lev(k, phs) result(res)

		integer, intent(in) :: k
		real(r8), intent(in) :: phs

		res = hyai(k) * p0 + hybi(k) * phs

	end function hybrid_coord_calc_ph_lev

	subroutine hybrid_coord_dcmip1(p0, hyai, hybi)

		real(r8), intent(in) :: p0
		real(r8), intent(out) :: hyai(num_lev+1)
		real(r8), intent(out) :: hybi(num_lev+1)

		real(r8), parameter :: ztop = 12000._r8
		real(r8), parameter :: T0   = 300._r8
		real(r8), parameter :: eta_top = exp(-g * ztop / Rd / T0)

		real(r8) dz, eta, z
		integer ilev

		dz = ztop / num_lev
		do ilev = 1, num_lev+1
			z = ztop - (ilev - 1) * dz
			eta = exp(-g * z / Rd / T0)
			hybi(ilev) = (eta - eta_top) / (1.0 - eta_top)
			hyai(ilev) = eta - hybi(ilev)
		end do

	end subroutine hybrid_coord_dcmip1

end module vert_coord_mod